package middleware

import (
	"bytes"
	"encoding/json"
	"fmt"
	"github.com/Shopify/sarama"
	"github.com/gin-gonic/gin"
	"github.com/olivere/elastic/v7"
	"github.com/sirupsen/logrus"
	"gopkg.in/sohlich/elogrus.v7"
	"io/ioutil"
	"log"
	"os"
	"time"
)

type responseBodyWriter struct {
	gin.ResponseWriter
	body *bytes.Buffer
}

func (w responseBodyWriter) Write(b []byte) (int, error) {
	w.body.Write(b)
	return w.ResponseWriter.Write(b)
}

func Logger(kafkaProducer sarama.SyncProducer) gin.HandlerFunc {
	logger := logrus.New()

	// 创建 Elastic Client
	options := []elastic.ClientOptionFunc{
		elastic.SetSniff(false),
		elastic.SetURL("http://localhost:9200"),
	}

	esclient, err := elastic.NewClient(options...)
	if err != nil {
		log.Panic(err)
	}

	// 创建 Elasticsearch Hook
	esHook, err := elogrus.NewElasticHook(esclient, "localhost", logrus.DebugLevel, "my_gin_blog")
	if err != nil {
		fmt.Println("Failed to create Elasticsearch hook:", err)
	}

	logger.AddHook(esHook)

	return func(c *gin.Context) {
		startTime := time.Now()

		// 使用自定义 ResponseWriter 来捕获响应体
		bodyWriter := &responseBodyWriter{body: bytes.NewBufferString(""), ResponseWriter: c.Writer}
		c.Writer = bodyWriter

		c.Next()

		stopTime := time.Since(startTime).Milliseconds()
		spendTime := fmt.Sprintf("%d ms", stopTime)
		hostName, err := os.Hostname()
		if err != nil {
			hostName = "unknown"
		}
		statusCode := c.Writer.Status()
		clientIp := c.ClientIP() //转为 请求端地址
		userAgent := c.Request.UserAgent()
		dataSize := c.Writer.Size()
		if dataSize < 0 {
			dataSize = 0
		}
		method := c.Request.Method
		path := c.Request.RequestURI

		// 请求日志信息
		requestLog := struct {
			ClientIP   string      `json:"client_ip"`
			UserAgent  string      `json:"user_agent"`
			Method     string      `json:"method"`
			Path       string      `json:"path"`
			Body       interface{} `json:"body"`
			StartTime  time.Time   `json:"start_time"`
			SpendTime  string      `json:"spend_time"`
			StatusCode int         `json:"status_code"`
		}{
			ClientIP:   clientIp,
			UserAgent:  userAgent,
			Method:     method,
			Path:       path,
			Body:       parseRequestBody(c),
			StartTime:  startTime,
			SpendTime:  spendTime,
			StatusCode: statusCode,
		}

		requestLogJSON, _ := json.Marshal(requestLog)

		// 返回日志信息
		responseLog := struct {
			StatusCode int         `json:"status_code"`
			Body       interface{} `json:"body"`
		}{
			StatusCode: statusCode,
			Body:       parseResponseBody(bodyWriter.body),
		}

		responseLogJSON, _ := json.Marshal(responseLog)

		entry := logger.WithFields(logrus.Fields{
			"HostName":  hostName,
			"status":    statusCode,
			"SpendTime": spendTime,
			"Ip":        clientIp,
			"Method":    method,
			"Path":      path,
			"DataSize":  dataSize,
			"Agent":     userAgent,
		})

		entry.Info("Request:", string(requestLogJSON), "Response:", string(responseLogJSON))

		// 将日志发送到 Kafka
		logEntry := map[string]interface{}{
			"HostName":  hostName,
			"status":    statusCode,
			"SpendTime": spendTime,
			"Ip":        clientIp,
			"Method":    method,
			"Path":      path,
			"DataSize":  dataSize,
			"Agent":     userAgent,
			"Request":   string(requestLogJSON),
			"Response":  string(responseLogJSON),
		}

		logEntryJSON, _ := json.Marshal(logEntry)
		message := &sarama.ProducerMessage{
			Topic: "my_gin_blog_logs",
			Value: sarama.StringEncoder(logEntryJSON),
		}

		_, _, err = kafkaProducer.SendMessage(message)
		if err != nil {
			fmt.Println("Failed to send log to Kafka:", err)
		}
	}
}

// ... 省略了 parseRequestBody 和 parseResponseBody 函数，不做修改 ...

// 辅助函数，解析请求体
func parseRequestBody(c *gin.Context) interface{} {
	body, err := ioutil.ReadAll(c.Request.Body)
	if err != nil {
		return nil
	}

	// 将请求体读取后，需要将其重新放入 Request.Body 中，以便后续中间件和处理器能够正常读取
	c.Request.Body = ioutil.NopCloser(bytes.NewBuffer(body))

	// 这里可以根据需要进行进一步的处理，比如解析 JSON 或其他格式的请求体
	var requestBody interface{}
	err = json.Unmarshal(body, &requestBody)
	if err != nil {
		return string(body) // 如果无法解析为 JSON，则返回字符串形式的请求体
	}

	return requestBody
}

// 辅助函数，解析响应体
func parseResponseBody(buffer *bytes.Buffer) interface{} {
	// 这里可以根据需要进行进一步的处理，比如解析 JSON 或其他格式的响应体
	var responseBody interface{}
	err := json.Unmarshal(buffer.Bytes(), &responseBody)
	if err != nil {
		return buffer.String() // 如果无法解析为 JSON，则返回字符串形式的响应体
	}

	return responseBody
}
